import os
import re
from datetime import datetime
from src.core.database import DatabaseConnector
from src.models.domain import Domain
from src.io.file_exporter import FileExporter
from src.io.csv_importer import CsvImporter
from src.core.help_generator import HelpGenerator
from src.core.dbhelper import DBHelper
import socket

class DomainCommand:
    """域名命令处理类"""
    
    def __init__(self, db_connector=None, subparsers=None):
        """初始化域名命令处理器"""
        self.db_connector = db_connector or DatabaseConnector()
        self.file_exporter = FileExporter()
        self.csv_importer = CsvImporter(self.db_connector)
        self.db_helper = DBHelper(self.db_connector)
        self.arg_parser = self.setup_parser(subparsers)

    def setup_parser(self, subparsers):
        """设置域名管理命令行参数解析器"""
        # domain命令组
        domain_parser = subparsers.add_parser('domain', help='域名管理相关命令', description='域名管理相关的所有操作')
        domain_subparsers = domain_parser.add_subparsers(dest='domain_command', help='域名管理子命令')

        # domain help命令组
        domain_subparsers.add_parser('help', help='显示域名管理命令帮助')

        # domain add命令组
        domain_add_parser = domain_subparsers.add_parser('add', help='添加域名')
        domain_add_parser.add_argument('domain', help='域名')
        domain_add_parser.add_argument('--owner', help='所属单位ID/名称')

        # domain set命令组
        domain_set_parser = domain_subparsers.add_parser('set', help='设置域名信息')
        domain_set_parser.add_argument('identifier', help='域名ID或域名')
        domain_set_parser.add_argument('--owner', help='所属单位ID/名称')

        # domain get命令组
        domain_get_parser = domain_subparsers.add_parser('get', help='获取域名信息')
        domain_get_parser.add_argument('identifier', help='域名ID或域名')
        domain_get_parser.add_argument('relation_type', nargs='?', choices=['ip', 'asset'], help='关系类型')

        # domain getall命令组
        domain_getall_parser = domain_subparsers.add_parser('getall', help='获取所有域名列表')
        domain_getall_parser.add_argument('-o', '--output', nargs='?', const=True, help='导出文件路径')
        domain_getall_parser.add_argument('-e', '--extension', choices=['json', 'csv'], default='csv', help='导出文件格式，默认csv')

        # domain search命令组
        domain_search_parser = domain_subparsers.add_parser('search', help='搜索域名')
        domain_search_parser.add_argument('keyword', help='搜索关键词')
        domain_search_parser.add_argument('-o', '--output', nargs='?', const=True, help='导出文件路径')
        domain_search_parser.add_argument('-e', '--extension', choices=['json', 'csv'], default='csv', help='导出文件格式，默认csv')


        # domain delete命令组
        domain_delete_parser = domain_subparsers.add_parser('delete', help='删除域名')
        domain_delete_parser.add_argument('identifier', help='域名ID或域名')

        # domain import命令组   
        domain_import_parser = domain_subparsers.add_parser('import', help='导入域名数据')
        domain_import_parser.add_argument('file_path', help='CSV文件路径')
        domain_import_parser.add_argument('--domain-column', help='域名列的索引或名称，默认为列名称为"domain"的列，或第一列', default='domain')
        domain_import_parser.add_argument('--owner-column', help='单位名称列的索引或名称，默认为列名称为"owner"的列，或第二列', default='owner')
        domain_import_parser.add_argument('--owner', help='所属单位ID/名称')
        domain_import_parser.add_argument('--skip-header', action='store_true', help='跳过CSV文件的第一行（标题行）')

    
        # domain check命令组
        domain_check_parser = domain_subparsers.add_parser('check', help='检查域名解析记录')
        domain_check_parser.add_argument('domain', help='域名')
        domain_check_parser.add_argument('--dns', help='DNS服务器IP地址，默认使用系统配置')
        domain_checkbatch_parser = domain_subparsers.add_parser('checkbatch', help='批量检查域名解析记录')
        domain_checkbatch_parser.add_argument('--check-filter', help='检查目标', choices=['all', 'invalid'], default='invalid')
        domain_checkbatch_parser.add_argument('--dns', help='DNS服务器IP地址，默认使用系统配置')

        return domain_parser
    
    def handle_command(self, args):
        """处理域名命令"""
        command = args.domain_command
        if command == 'help':
            self.show_domain_help()
        elif command == 'add':
            self.add_domain(args.domain, args.owner)
        elif command == 'set':
            self.set_domain(args.identifier, args.owner)
        elif command == 'get':
            self.get_domain(args.identifier, args.relation_type)
        elif command == 'getall':
            export_path = args.output if hasattr(args, 'export') else None
            extension = args.extension if hasattr(args, 'extension') else None
            self.get_all_domains(export_path=export_path, extension=extension)
        elif command == 'search':
            export_path = args.output if hasattr(args, 'export') else None
            extension = args.extension if hasattr(args, 'extension') else None
            self.search_domains(args.keyword, export_path, extension)
        elif command == 'delete':
            self.delete_domain(args.identifier)
        elif command == 'import':
            self.import_domains(args.file_path, args.domain_column, args.owner_column, args.owner)
        elif command == 'check':
            self.check_domain(args.domain, args.dns)
        elif command == 'checkbatch':
            self.check_domains(args.check_filter, args.dns)
        else:
            print(f"未知的域名命令: {command}")
            self.show_domain_help()
    
    def show_domain_help(self):
        """显示域名命令的帮助信息"""
        self.arg_parser.print_help()
    
    def is_valid_domain(self, domain):
        """验证域名格式是否有效"""
        # 简单的域名验证正则表达式
        pattern = r'^([a-zA-Z0-9]([a-zA-Z0-9\-]{0,61}[a-zA-Z0-9])?\.)+[a-zA-Z]{2,}$'
        return re.match(pattern, domain) is not None
    
    def _get_owner(self, identifier):
        """获取单位"""
        return self.db_helper.get_owner(identifier)
           
    def add_domain(self, domain_name, owner_identifier=None):
        """添加新域名"""
        domain = self._get_domain(domain_name)
        if domain:
            return
        
        print(f"添加域名: {domain_name}")
        
        owner_id = None
        if owner_identifier:
            owner = self._get_owner(owner_identifier)
            if not owner:
                print(f"未找到单位: {owner_identifier}")
                return
            owner_id = owner['owner_id']
        # 添加域名
        current_time = datetime.now()
        sql = "INSERT INTO domains (domain, owner_id, create_time, update_time) VALUES (%s, %s, %s, %s)"
        try:
            cursor = self.db_connector.execute_query(sql, (domain_name, owner_id, current_time, current_time))
            if cursor:
                domain_id = cursor.lastrowid if hasattr(cursor, 'lastrowid') else '未知'
                print(f"成功添加域名: ID={domain_id}")
                print(f"  域名: {domain_name}")
                if owner_id:
                    owner_name = owner['owner_name']
                    print(f"  所属单位: {owner_name} (ID: {owner_id})")
        except Exception as e:
            print(f"添加域名失败: {e}")
    
    def _get_domain(self, identifier):
        """获取指定域名信息"""
        return self.db_helper.get_domain(identifier)

    def set_domain(self, identifier, owner_identifier=None):
        """设置域名所属单位"""
        domain = self._get_domain(identifier)
        if not domain:
            print(f"未找到域名: {identifier}")
            return

        if owner_identifier:
            owner = self._get_owner(owner_identifier) if owner_identifier else None
            if not owner:
                print(f"错误：未找到单位: {owner_identifier}")
                return
            owner_id = owner['owner_id']
            owner_name = owner['owner_name']

            # 更新域名所属单位
            current_time = datetime.now()
            sql = "UPDATE domains SET owner_id = %s, update_time = %s WHERE domain_id = %s"
            params = [owner_id, current_time, domain['domain_id']]
            self.db_connector.execute_query(sql, params)
            print(f"成功设置域名 {domain['domain']} 的所属单位为 {owner_name} (ID: {owner_id})")

    def get_domain(self, identifier, relation_type):
        """获取指定域名信息"""
        domain = self._get_domain(identifier)
        if not domain:
            print(f"未找到域名: {identifier}")
            return None

        print("域名信息:")
        print(f"  ID: {domain['domain_id']}")
        print(f"  域名: {domain['domain']}")
        if domain['owner_id']:
            print(f"  所属单位: {domain['owner_name']} (ID: {domain['owner_id']})")
        else:
            print("  所属单位: 无")
        print(f"  创建时间: {domain['create_time']}")
        print(f"  更新时间: {domain['update_time']}")
        
        print("")
        if relation_type == 'asset':
            self._show_domain_assets(domain['domain'])
        elif relation_type == 'ip':
            self._show_domain_ips(domain['domain'])
    
    def get_all_domains(self, export_path=None, extension=None):
        """获取所有域名列表"""
        # 查询所有域名
        sql = "SELECT d.*, o.owner_name FROM domains d LEFT JOIN owners o ON d.owner_id = o.owner_id ORDER BY d.domain"
        domains_result = self.db_connector.fetchall(sql)
        
        if len(domains_result) == 0:
            print("暂无域名信息！")
            return
        
        # 准备导出数据
        export_data = []
        for domain in domains_result:
            export_item = {
                'domain_id': domain['domain_id'],
                'domain': domain['domain'],
                'owner_id': domain['owner_id'],
                'create_time': domain['create_time'],
                'update_time': domain['update_time'],
                'owner_name': domain['owner_name']
            }
            # 转换datetime对象为字符串
            if 'create_time' in export_item and isinstance(export_item['create_time'], datetime):
                export_item['create_time'] = export_item['create_time'].strftime('%Y-%m-%d %H:%M:%S')
            if 'update_time' in export_item and isinstance(export_item['update_time'], datetime):
                export_item['update_time'] = export_item['update_time'].strftime('%Y-%m-%d %H:%M:%S')
            export_data.append(export_item)
        
        # 导出数据
        if export_path is not None:
            actual_path = None if export_path is True else export_path
            self.file_exporter.export_data(export_data, actual_path, 'all_domains', extension)
            return
        
        # 显示域名列表
        print(f"共有 {len(domains_result)} 个域名:")
        print("{:<10} {:<30} {:<20} {:<20} {:<20}".format(
            "ID", "域名", "所属单位", "创建时间", "更新时间"
        ))
        print("=" * 100)
        
        for row in domains_result:
            domain_id = row['domain_id']
            domain_name = row['domain']
            owner_id = row['owner_id']
            owner_name = row['owner_name']
            create_time = row['create_time']
            update_time = row['update_time']
            
            # 转换datetime对象为字符串
            if isinstance(create_time, datetime):
                create_time = create_time.strftime('%Y-%m-%d %H:%M:%S')
            if isinstance(update_time, datetime):
                update_time = update_time.strftime('%Y-%m-%d %H:%M:%S')
            
            print("{:<10} {:<30} {:<20} {:<20} {:<20}".format(
                domain_id, domain_name, f"{owner_name} (ID: {owner_id})", create_time, update_time
            ))
            
        print(f"共有 {len(domains_result)} 个域名")
    
    def search_domains(self, keyword, export_path=None, extension=None):
        """搜索域名"""
        # 模糊匹配域名或所属单位名称
        sql = "SELECT d.*, o.owner_name FROM domains d LEFT JOIN owners o ON d.owner_id = o.owner_id WHERE d.domain LIKE %s OR o.owner_name LIKE %s ORDER BY d.domain"
        params = [f"%{keyword}%", f"%{keyword}%"]
        
        domains_result = self.db_connector.fetchall(sql, params)
        
        if len(domains_result) == 0:
            print(f"未找到包含 '{keyword}' 的域名！")
            return
        
        # 准备导出数据
        export_data = []
        for domain in domains_result:
            export_item = {
                'domain_id': domain['domain_id'],
                'domain': domain['domain'],
                'owner_id': domain['owner_id'],
                'create_time': domain['create_time'],
                'update_time': domain['update_time'],
                'owner_name': domain['owner_name']
            }
            # 转换datetime对象为字符串
            if 'create_time' in export_item and isinstance(export_item['create_time'], datetime):
                export_item['create_time'] = export_item['create_time'].strftime('%Y-%m-%d %H:%M:%S')
            if 'update_time' in export_item and isinstance(export_item['update_time'], datetime):
                export_item['update_time'] = export_item['update_time'].strftime('%Y-%m-%d %H:%M:%S')
            export_data.append(export_item)
        
        # 导出数据
        if export_path is not None:
            actual_path = None if export_path is True else export_path
            self.file_exporter.export_data(export_data, actual_path, 'domain_search', extension)
            return
        
        # 显示搜索结果
        print(f"找到 {len(domains_result)} 个匹配的域名:")
        print("{:<10} {:<30} {:<20} {:<20} {:<20}".format(
            "ID", "域名", "所属单位", "创建时间", "更新时间"
        ))
        print("=" * 100)
        
        for row in domains_result:
            domain_id = row['domain_id']
            domain_name = row['domain']
            owner_id = row['owner_id']
            owner_name = row['owner_name']
            create_time = row['create_time']
            update_time = row['update_time']
            
            # 转换datetime对象为字符串
            if isinstance(create_time, datetime):
                create_time = create_time.strftime('%Y-%m-%d %H:%M:%S')
            if isinstance(update_time, datetime):
                update_time = update_time.strftime('%Y-%m-%d %H:%M:%S')
            
            print("{:<10} {:<30} {:<20} {:<20} {:<20}".format(
                domain_id, domain_name, f"{owner_name} (ID: {owner_id})", create_time, update_time
            ))
    
    def delete_domain(self, identifier):
        """删除域名"""
        domain = self._get_domain(identifier)
        
        if not domain:
            print(f"未找到域名: {identifier}")
            return
        
        domain_id = domain['domain_id']
        # 删除域名
        sql = "DELETE FROM domains WHERE domain_id = %s"
        try:
            self.db_connector.execute_query(sql, (domain_id,))
            print(f"成功删除域名: {domain['domain']} (ID={domain_id})")
        except Exception as e:
            print(f"删除域名失败: {e}")
    
    def import_domains(self, file_path, domain_column = 'domain', owner_column = 'owner', owner_identifier=None):
        """从CSV文件导入域名数据"""
        # 验证文件是否存在
        if not os.path.exists(file_path):
            print(f"错误: 文件 {file_path} 不存在！")
            return
        
        # 检查文件扩展名
        file_ext = os.path.splitext(file_path)[1].lower()
        if file_ext != '.csv':
            print(f"错误: 只支持CSV文件格式，当前文件格式为 '{file_ext}'")
            return
        
        try:
            owner_id = None
            owner = self.db_helper.get_or_add_owner(owner_identifier)
            if owner:
                owner_id = owner['owner_id']

            success_count, failed_count = self.csv_importer.import_domains_from_csv(file_path, domain_column, owner_column, owner_id)
            print(f"域名导入完成: 成功 {success_count} 个, 失败 {failed_count} 个")
        except Exception as e:
            print(f"导入域名时发生错误: {str(e)}")
    
    def _show_domain_assets(self, domain_name):
        """显示域名关联的资产信息"""
        # 查询关联的资产数量
        sql = "SELECT * FROM assets WHERE domain = %s"
        assets_result = self.db_connector.fetchall(sql, (domain_name,))
        
        if len(assets_result) == 0:
            print(f"域名 '{domain_name}' 下没有的资产！")
            return

        print(f"域名 '{domain_name}' 资产信息:")
        print("{:<10} {:<30} {:<20} {:<15} {:<5} {:<3}".format(
            "ID", "名称", "域名", "IP", "端口", "协议"
        ))
        print("=" * 100)
        for asset in assets_result:
            ip = asset['ip'] if asset['ip'] is not None else ''
            domain = asset['domain'] if asset['domain'] is not None else ''
            print("{:<10} {:<30} {:<20} {:<15} {:<5} {:<3}".format(
                asset['asset_id'], asset['asset_name'], domain, ip, asset['port'], asset['asset_protocol']
            ))
        
    def _show_domain_ips(self, domain_name):
        """显示域名关联的IP地址信息"""
        sql = "SELECT * FROM ipdomains ido LEFT JOIN ips on ido.ip = ips.ip WHERE domain = %s"
        ips_result = self.db_connector.fetchall(sql, (domain_name,))

        if len(ips_result) == 0:
            print(f"域名 '{domain_name}' 下没有关联的IP地址！")
            return

        print(f"域名 '{domain_name}' 关联的IP地址信息:")
        print("{:<10} {:<20} {:<10} {:<20}".format("ID", "IP地址", "是否有效", "最后DNS解析时间"))
        print("=" * 65)

        for ip in ips_result:
            ip_id = ip['ip_id']
            ip_address = ip['ip']
            is_valid = ip['is_valid']
            last_check_time = ip['last_check_time']
            if last_check_time is not None:
                last_check_time = last_check_time.strftime('%Y-%m-%d %H:%M:%S')
            else:
                last_check_time = 'N/A'
            print("{:<10} {:<20} {:<10} {:<20}".format(ip_id, ip_address, '有效' if is_valid else '无效', last_check_time))

    def check_domain(self, domain_name, dns=None):
        """检查域名是否存在"""
        domain = self.db_helper.get_domain(domain_name)
        
        # 开始通过DNS解析域名
        ips = self._resolve_domain(domain_name, dns)
        if len(ips) > 0:
            print(f"域名 '{domain_name}' 解析成功，IP地址为: {ips}")
            for ip in ips:
                try:
                    self.db_helper.check_ipdomain(ip, domain_name, True)
                except Exception as e:
                    print(f"添加IP和域名关联失败: {e}")
        else:
            print(f"域名 '{domain_name}' 解析失败！")

    def check_domains(self, check_filter, dns=None):
        sql = ''
        if check_filter == 'all':
            sql = "SELECT * FROM domains"
        elif check_filter == 'invalid':
            sql = "SELECT * FROM domains WHERE domain NOT IN (SELECT domain FROM ipdomains)"
        else:
            print(f"无效的检查过滤条件: {check_filter}")
            return

        domains_result = self.db_connector.fetchall(sql)

        if len(domains_result) == 0:
            print(f"没有符合条件的域名！")
            return

        for domain in domains_result:
            self.check_domain(domain['domain'], dns)
            print("")

    def _resolve_domain(self, domain_name, dns=None):
        """通过DNS解析域名"""
        try:
            if dns is None:
                ips = socket.gethostbyname_ex(domain_name)[2]
            else:
                # 创建一个UDP套接字
                sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
                # 设置超时时间为2秒
                sock.settimeout(2)
                # DNS查询数据包
                query_id = b'\x01\x00'  # ID
                query_flags = b'\x01\x00'  # 标准查询
                query_qdcount = b'\x00\x01'  # 1个问题
                query_ancount = b'\x00\x00'  # 0个答案
                query_nscount = b'\x00\x00'  # 0个权威记录
                query_arcount = b'\x00\x00'  # 0个额外记录
                
                # 构建查询问题部分
                qname = b''
                for part in domain_name.split('.'):
                    qname += bytes([len(part)]) + part.encode()
                qname += b'\x00'  # 结束标记
                qtype = b'\x00\x01'  # A记录
                qclass = b'\x00\x01'  # IN
                
                # 组合完整的DNS查询数据包
                dns_query = (query_id + query_flags + query_qdcount + query_ancount + 
                            query_nscount + query_arcount + qname + qtype + qclass)
                
                # 发送DNS查询
                sock.sendto(dns_query, (dns, 53))
                
                # 接收DNS响应
                data, _ = sock.recvfrom(512)  # DNS响应最大512字节
                sock.close()
                
                # 解析响应中的IP地址（简化版）
                ips = []
                # 跳过头部和问题部分
                ptr = 12  # 头部长度
                # 跳过问题部分
                while data[ptr] != 0:
                    ptr += data[ptr] + 1
                ptr += 5  # 跳过0结束符和qtype、qclass
                
                # 解析答案部分
                ancount = int.from_bytes(data[6:8], byteorder='big')
                for _ in range(ancount):
                    # 跳过域名部分（使用压缩格式）
                    if (data[ptr] & 0xc0) == 0xc0:
                        ptr += 2
                    else:
                        while data[ptr] != 0:
                            ptr += data[ptr] + 1
                        ptr += 1
                    
                    # 跳过类型和类
                    ptr += 4
                    # 跳过生存时间
                    ptr += 4
                    # 获取数据长度
                    data_len = int.from_bytes(data[ptr:ptr+2], byteorder='big')
                    ptr += 2
                    
                    # 解析IPv4地址
                    if data_len == 4:
                        ip = '.'.join(str(b) for b in data[ptr:ptr+4])
                        ips.append(ip)
                    ptr += data_len
                
                # 如果上面的自定义DNS解析失败，回退到系统解析
                if not ips:
                    ips = socket.gethostbyname_ex(domain_name)[2]
            return ips
        except (socket.gaierror, socket.timeout, Exception):
            # print(f"域名解析失败: {domain_name}")
            return []
